"use client";

import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import { v4 as uuid } from "uuid";
import { openStream, requestDevToken } from "@/lib/api";
import { parseStream } from "@/lib/streamParser";
import { GraphExtractor } from "@/lib/graphExtractor";
import type {
  ChatMessage,
  ContextGraph,
  ToolInvocation,
  ToolSelection,
} from "@/lib/types";

interface UseChatStreamReturn {
  messages: ChatMessage[];
  graph: ContextGraph;
  setGraph: (g: ContextGraph) => void;
  extractor: GraphExtractor;
  pending: boolean;
  conversationId: string;
  tokenReady: boolean;
  error: string | null;
  send: (prompt: string, tools: ToolSelection[]) => Promise<void>;
  cancel: () => void;
  resetConversation: () => void;
}

export function useChatStream(): UseChatStreamReturn {
  const [messages, setMessages] = useState<ChatMessage[]>([]);
  const [graph, setGraph] = useState<ContextGraph>({ nodes: [], rels: [] });
  const [pending, setPending] = useState(false);
  const [conversationId, setConversationId] = useState(() => uuid());
  const [tokenReady, setTokenReady] = useState(false);
  const [error, setError] = useState<string | null>(null);

  const extractor = useMemo(() => new GraphExtractor(), []);
  const abortRef = useRef<AbortController | null>(null);

  // Bootstrap: fetch a dev JWT once on mount
  useEffect(() => {
    let cancelled = false;
    (async () => {
      try {
        await requestDevToken();
        if (!cancelled) setTokenReady(true);
      } catch (e) {
        if (!cancelled) setError((e as Error).message);
      }
    })();
    return () => {
      cancelled = true;
    };
  }, []);

  const send = useCallback(
    async (prompt: string, tools: ToolSelection[]) => {
      if (!prompt.trim()) return;
      setError(null);
      setPending(true);

      // Pre-flight: verify backend + dependencies are healthy before burning
      // an LLM turn on a request that would 500 anyway.
      try {
        const probe = await fetch("/api/health", { cache: "no-store" });
        if (!probe.ok) {
          setError(`Service not ready (HTTP ${probe.status}). Check backend status.`);
          setPending(false);
          return;
        }
      } catch (e) {
        setError(`Backend unreachable: ${(e as Error).message}`);
        setPending(false);
        return;
      }

      // Wipe the previous prompt's graph so each question gets a clean canvas.
      // Evidence of prior turns still lives in the chat trace; the graph
      // visualises the *current* reasoning path.
      extractor.reset();
      setGraph({ nodes: [], rels: [] });

      const userMsg: ChatMessage = {
        id: uuid(),
        role: "user",
        content: prompt,
        createdAt: Date.now(),
      };
      const asstMsg: ChatMessage = {
        id: uuid(),
        role: "assistant",
        content: "",
        createdAt: Date.now(),
        toolCalls: [],
      };
      setMessages((m) => [...m, userMsg, asstMsg]);

      const toolInvocations = new Map<string, ToolInvocation>();
      const ac = new AbortController();
      abortRef.current = ac;

      try {
        const res = await openStream(
          {
            prompt,
            conversation_id: conversationId,
            include_tools: tools,
          },
          ac.signal
        );

        for await (const ev of parseStream(res)) {
          if (ev.type === "content" && ev.content) {
            const text = ev.content;
            setMessages((m) =>
              m.map((msg) =>
                msg.id === asstMsg.id
                  ? { ...msg, content: msg.content + text }
                  : msg
              )
            );
          } else if (ev.type === "tool_call" && ev.tool_id) {
            const inv: ToolInvocation = {
              toolId: ev.tool_id,
              tool: ev.tool ?? "unknown",
              input: ev.tool_input ?? {},
              startedAt: Date.now(),
              status: "pending",
            };
            toolInvocations.set(ev.tool_id, inv);
            setMessages((m) =>
              m.map((msg) =>
                msg.id === asstMsg.id
                  ? {
                      ...msg,
                      toolCalls: [...(msg.toolCalls ?? []), inv],
                    }
                  : msg
              )
            );
            setGraph(extractor.ingest(inv));
          } else if (ev.type === "tool_result" && ev.tool_id) {
            const inv = toolInvocations.get(ev.tool_id);
            if (inv) {
              inv.output = ev.tool_output;
              inv.endedAt = Date.now();
              inv.status = "ok";
              setMessages((m) =>
                m.map((msg) =>
                  msg.id === asstMsg.id
                    ? {
                        ...msg,
                        toolCalls: (msg.toolCalls ?? []).map((t) =>
                          t.toolId === inv.toolId ? { ...inv } : t
                        ),
                      }
                    : msg
                )
              );
              setGraph(extractor.ingest(inv));
            }
          } else if (ev.type === "error" && ev.content) {
            setError(ev.content);
          }
        }
      } catch (e) {
        if ((e as Error).name !== "AbortError") {
          setError((e as Error).message);
        }
      } finally {
        setPending(false);
        abortRef.current = null;
      }
    },
    [conversationId, extractor]
  );

  const cancel = useCallback(() => {
    abortRef.current?.abort();
  }, []);

  const resetConversation = useCallback(() => {
    cancel();
    setMessages([]);
    setGraph({ nodes: [], rels: [] });
    extractor.reset();
    setConversationId(uuid());
    setError(null);
  }, [cancel, extractor]);

  return {
    messages,
    graph,
    setGraph,
    extractor,
    pending,
    conversationId,
    tokenReady,
    error,
    send,
    cancel,
    resetConversation,
  };
}
